import pandas as pd
import numpy as np
from typing import Union, List
import sys
from contextlib import contextmanager
import os

@contextmanager
def suppress_stdout():
    with open(os.devnull, 'w') as devnull:
        old_stdout = sys.stdout
        sys.stdout = devnull
        sys.stderr = devnull
        try:
            yield
        finally:
            sys.stdout = old_stdout



def pre_data(u_nk1: 'pd.DataFrame', u_nk2: 'pd.DataFrame', n_k1: 'np.ndarray', n_k2: 'np.ndarray',
             lambdas: Union['np.ndarray', List]):
    nstates = len(lambdas) // 2
    u_nk1['fep-lambda'] = lambdas[:nstates]
    u_nk1.set_index('fep-lambda', inplace=True)
    u_nk1 = u_nk1.T
    tmp = [[val, ] * i for val, i in zip(lambdas, n_k1)]
    u_nk1['fep-lambda'] = [j for i in tmp for j in i]
    u_nk1['times'] = [i for i in range(u_nk1.shape[0])]
    u_nk1.set_index(['times', 'fep-lambda'], inplace=True)

    u_nk2['fep-lambda'] = lambdas[11:]
    u_nk2.set_index('fep-lambda', inplace=True)
    u_nk2 = u_nk2.T
    tmp = [[val, ] * i for val, i in zip(lambdas[11:], n_k2)]
    u_nk2['fep-lambda'] = [j for i in tmp for j in i]
    u_nk2['times'] = [i for i in range(u_nk2.shape[0])]
    u_nk2.set_index(['times', 'fep-lambda'], inplace=True)

    return u_nk1, u_nk2


def fep(u_nk1: 'pd.DataFrame', u_nk2: 'pd.DataFrame'):
    with suppress_stdout():
        from alchemlyb.estimators import MBAR
        mbar_lambda1 = MBAR()
        mbar_lambda1.fit(u_nk1)
        mbar_lambda2 = MBAR()
        mbar_lambda2.fit(u_nk2)
        return mbar_lambda1.delta_f_.loc[0.00, 0.50] + mbar_lambda2.delta_f_.loc[0.50, 1.00]

def time_convergence(u_nk1: 'pd.DataFrame', u_nk2: 'pd.DataFrame'):
    with suppress_stdout():
        ## 评估收敛
        from alchemlyb.convergence import forward_backward_convergence
        from alchemlyb.visualisation import plot_convergence
        u_nk1_groups = u_nk1.groupby(level=u_nk1.index.names[1:])
        df = forward_backward_convergence([i for _, i in u_nk1_groups])
        ax = plot_convergence(df)
        ax.figure.savefig('dF_t_nk1.pdf')

        u_nk2_groups = u_nk2.groupby(level=u_nk2.index.names[1:])
        df = forward_backward_convergence([i for _, i in u_nk2_groups])
        ax = plot_convergence(df)
        ax.figure.savefig('dF_t_nk2.pdf')


def block_average(u_nk1: 'pd.DataFrame', u_nk2: 'pd.DataFrame'):
    with suppress_stdout():
        from alchemlyb.visualisation import plot_block_average
        from alchemlyb.convergence import block_average
        u_nk1_groups = u_nk1.groupby(level=u_nk1.index.names[1:])
        df1 = block_average([i for _, i in u_nk1_groups])
        ax1 = plot_block_average(df1)
        ax1.figure.savefig('dF1_t_block_average.png')

        u_nk2_groups = u_nk2.groupby(level=u_nk2.index.names[1:])
        df2 = block_average([i for _, i in u_nk2_groups])
        ax2 = plot_block_average(df2)
        ax2.figure.savefig('dF2_t_block_average.png')

def overlap_matrix():
    pass
